import copy
import random
from typing import Any, Dict, Optional
import numpy as np
import gymnasium as gym
from gymnasium import spaces
import cv2
import platform

from gym_pih.hardware.RealsenseCamera import Camera as RealsenseCamera
from gym_pih.hardware.ft300 import RobotiqFt300 as FT300
from gym_pih.hardware.gpio import GPIO
from gym_pih.hardware.pid import PIDForceController
from gym_pih.hardware.calculate_slope import calculate_slope
from gym_pih.hardware.LightSource import LightSource
from gym_pih.reward_classifier.load_classifier import Classifier
import rtde_control
import rtde_receive

if platform.system() == 'Windows':
    from gym_pih.hardware.Imagingsource_for_windows import Camera as ImagingsourceCamera
elif platform.system() == 'Linux':
    from gym_pih.hardware.Imagingsource_for_linux import Camera as ImagingsourceCamera
else:
    raise Exception("Unsupported platform")

class TypecEnv(gym.Env):
    def __init__(self,
                 target_pos: np.ndarray = np.array(
                     [0.6872853048650824, -0.1308957273597933, 0.17394818223574388, 6.256728706533603e-05, -3.141092053034699, 0.0001309061484982937] ),
                 reward_shape="sparse",
                 train_env='normal_env',
                 fake_env=False
                 ):
        self.params_pos_error=0   #孔的定位误差范围
        self.action_range=0.001   #动作范围1mm
        self.initial_error=0.01   #初始位置与孔的最大距离
        self.ministep_num=6       #每次动作分8步执行
        
        self.train_env=train_env
        self.ip="192.168.3.140"
        self.velocity = 0.1
        self.acceleration = 0.05
        self.dt = 1.0/125  # 8ms
        self.lookahead_time = 0.1
        self.gain = 300
        
        assert reward_shape in ["sparse","dense"],f"reward_shape must be in ['sparse','dense'],but got {reward_shape}"
        self.reward_shape=reward_shape
        self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
        if train_env=='serl':
            self.observation_space = gym.spaces.Dict(
                {
                    "state": gym.spaces.Dict(
                        {
                            "position": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float32),
                            "force":gym.spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float32),
                            "slope":gym.spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32),
                        }
                    ),
                    "images": gym.spaces.Dict(
                        {
                            "wrist": gym.spaces.Box(
                                low=0,
                                high=255,
                                shape=(64, 64, 3),
                                dtype=np.uint8,
                            ),
                            "side": gym.spaces.Box(
                                low=0,
                                high=255,
                                shape=(64, 64, 3),
                                dtype=np.uint8,
                            ),
                        }
                    ),
                }
            )

        else:
            self.observation_space = gym.spaces.Dict(
                {"state":gym.spaces.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32),
                 "wrist":gym.spaces.Box(low=0,high=1,shape=(3,64,64),dtype=np.float32),
                 "side":gym.spaces.Box(low=0,high=1,shape=(3,64,64),dtype=np.float32),}
            )

        

        self.pre_distance=None
        self.pre_slope=0
        assert target_pos.shape==(6,),f"target_pos must be a 6-dim ndarray,but got {target_pos.shape}"
        self.target_pos=target_pos.copy() 
        self.pos_error=np.random.uniform(-self.params_pos_error,self.params_pos_error,2)
        self.target_pos[:2]+=self.pos_error#孔的位置有定位误差
        self.pose=np.zeros(6)#用于控制机器人运动
        self.boundary=np.zeros((2,3))#用于限制机器人的运动范围
        self.boundary[0]=self.target_pos[:3]-0.025
        self.boundary[0][2]=self.target_pos[2]-0.003
        self.boundary[1]=self.target_pos[:3]+0.025

        if fake_env:
            return
        self.Light=LightSource()
        self.Light.Light_Open(6)
        self.Realsense_cam = RealsenseCamera()
        self.Realsense_cam.Open()
        self.Imagingsource_cam = ImagingsourceCamera()
        self.Imagingsource_cam.Open()
        if platform.system() == 'Windows':
            self.Imagingsource_cam.SetProperty("Gain","Value",0)
            self.Imagingsource_cam.SetProperty("Exposure","Value",float(1/30))
        self.ft=FT300()
        # self.gpio=GPIO()
        # self.gpio.IO_InitPin(0,0,2)
        self.action_controller=PIDForceController(kp=self.action_range/self.ministep_num/4,
                                                  ki=self.action_range/self.ministep_num/100/150,
                                                  kd=self.action_range/self.ministep_num/4)
        self.rtde_c = rtde_control.RTDEControlInterface(self.ip)
        self.rtde_r = rtde_receive.RTDEReceiveInterface(self.ip)

        self.classifier=Classifier()

    def _get_obs(self):
        target_pos=self.target_pos.copy()[:3]
        pos_ee=self.rtde_r.getActualTCPPose()[:3]
        force=self.ft.forceTorque[:3]
        _,wrist_image=self.Realsense_cam.GetImage(isShow=False)
        _,side_image=self.Imagingsource_cam.GetImage()
        distance=pos_ee-target_pos
        pos_ee=np.array(pos_ee,dtype=np.float32)
        distance=np.array(distance,dtype=np.float32)
        force=np.array(force,dtype=np.float32)
        if self.train_env=='serl':
            wrist_image=wrist_image[150:214,210:274]
            done,slope=calculate_slope(side_image[130:,250:1020])
            if not done:
                slope=self.pre_slope
            else:
                self.pre_slope=slope
            slope=np.array([abs(slope)],dtype=np.float32)
            side_image=side_image[40:720, 260:900]
            side_image=cv2.resize(side_image,(64,64))
            observation={"state":{"position":pos_ee,"force":force,"slope":slope},
                         "images":{"wrist":wrist_image,"side":side_image}}
            observation=self._normalize(observation)
        else:
            wrist_image=wrist_image[150:214,210:274]
            done,slope=calculate_slope(side_image[130:,250:1020])
            if not done:
                slope=self.pre_slope
            else:
                self.pre_slope=slope
            slope=np.array([abs(slope)],dtype=np.float32)
            side_image=side_image[40:720, 260:900]
            side_image=cv2.resize(side_image,(64,64))
            state=np.concatenate((pos_ee,force,slope))
            wrist_image= np.transpose(wrist_image, (2, 0, 1))
            side_image = np.transpose(side_image, (2, 0, 1))
            observation={"state":state,"wrist":wrist_image,"side":side_image}
            observation=self._normalize(observation)
        return observation
    
    def _get_info(self):
        return {}
    
    def _get_reward(self,observation,terminated):
        if self.train_env=="serl":
            reward=-0.005
            insert_state=self.classifier.predict(observation)
            reward+=insert_state*0.002
            if insert_state:
                slope=(observation["state"]["slope"]).item()
                reward+=(slope-0.5)*0.002
            if terminated:
                reward=1
        else:
            reward = -0.005
            obs={"images":{"wrist":np.transpose(observation["wrist"],(1,2,0)),"side":np.transpose(observation["side"],(1,2,0))}}
            insert_state = self.classifier.predict(obs)
            reward += insert_state * 0.002
            if insert_state:
                slope = (observation["state"][6]).item()
                reward += (slope - 0.5) * 0.002
            if terminated:
                reward = 1
        return reward
    
    def _distance(self,distance_array):
        return np.linalg.norm(distance_array)
    
    def reset(self, seed: int = None, options: Optional[Dict[str, Any]] = None):
        """先向上移动拔出插头，再复位"""
        #super().reset(seed=seed, options=options)
        self.rtde_c.servoStop()        #ur进入servo模式后，必须先停止才能再次用move
        init_pose=self.target_pos.copy()
        init_pose[:2] += np.random.uniform(-self.initial_error,self.initial_error, 2)
        init_pose[2]+=0.015
        transition_pose=self.rtde_r.getActualTCPPose()
        # transition_pose[2]+=0.003
        # self.rtde_c.moveL(transition_pose, 0.2, 0.05)
        transition_pose[1]+=random.uniform(-0.05, -0.02) if init_pose[1]<self.target_pos[1] else random.uniform(0.01, 0.02)
        self.rtde_c.moveL(transition_pose, 0.2, 0.05)
        transition_pose[2]=init_pose[2]
        self.rtde_c.moveL(transition_pose, 0.2, 0.05)
        self.rtde_c.moveL(init_pose, 0.2, 0.05)

        self.pose=self.rtde_r.getActualTCPPose()
        self.pre_distance=None
        observation=self._get_obs()
        info=self._get_info()

        return observation,info
    
    def step(self, action):
        self._miniServoL(action*self.action_range,self.ministep_num)
        terminated=self.rtde_r.getDigitalInState(0)
        observation=self._get_obs()
        reward=self._get_reward(observation,terminated)
        info=self._get_info()

        return observation, reward, terminated, False, info

    def _miniServoL(self, action, step=4):
        """
        输入一次动作，分多步执行

        :param action: 动作
        :param step: 步数，默认为4
        """
        miniaction = action / step
        for i in range(step):
            force=self.ft.GetForce()
            delta=self.action_controller._calPIDForce(force)
            self.pose[:3] += miniaction+delta[:3]
            self.pose[:3]=np.clip(self.pose[:3],self.boundary[0],self.boundary[1])
            t_start=self.rtde_c.initPeriod()
            self.rtde_c.servoL(self.pose, self.velocity, self.acceleration, self.dt, self.lookahead_time, self.gain)
            self.rtde_c.waitPeriod(t_start)

    def _normalize(self,observation):
        """
        观测值归一化

        :param observation: 观测值
        :return: 归一化后的观测值
        """
        normal_obs = copy.deepcopy(observation)
        if self.train_env=='serl':
            normal_obs['state']['position']=((normal_obs['state']['position'] - self.boundary[0]) / (self.boundary[1] - self.boundary[0])).astype(np.float32)
            normal_obs['state']['force']=np.clip(normal_obs['state']['force'],-10,10)/10
            normal_obs['state']['slope']/=90
        else:
            normal_obs["state"][0:3]=((normal_obs["state"][0:3]-self.boundary[0]) / (self.boundary[1] - self.boundary[0])).astype(np.float32)
            normal_obs["state"][3:6]=np.clip(normal_obs["state"][3:6],-10,10)/10
            normal_obs["state"][6]/=90
            normal_obs["wrist"]=(normal_obs["wrist"]/255.0).astype(np.float32)
            normal_obs["side"]=(normal_obs["side"]/255.0).astype(np.float32)
        return normal_obs
    
    def close(self):
        self.rtde_c.servoStop()
        self.rtde_c.stopScript()
        self.rtde_c.disconnect()
        self.rtde_r.disconnect()
        self.ft.DisConnect()
        self.Light.Light_Close()
        self.Light.port_close()
